{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql import SparkSession\n",
    "from pyspark.ml.feature import HashingTF\n",
    "from pyspark.ml.classification import LogisticRegression\n",
    "from pyspark.ml.evaluation import BinaryClassificationEvaluator\n",
    "from pyspark.sql.functions import udf\n",
    "from pyspark.sql.types import ArrayType, StringType\n",
    "import jieba\n",
    "import time\n",
    "\n",
    "def initialize_spark():\n",
    "    return SparkSession.builder \\\n",
    "        .appName(\"TextClassification\") \\\n",
    "        .master(\"local[12]\") \\\n",
    "        .config(\"spark.executor.memory\", \"4g\") \\\n",
    "        .config(\"spark.executor.cores\", \"5\") \\\n",
    "        .config(\"spark.task.cpus\", \"1\") \\\n",
    "        .config(\"spark.executor.instances\", \"2\") \\\n",
    "        .config(\"spark.sql.auto.repartition\", \"true\") \\\n",
    "        .config(\"spark.driver.memory\", \"8g\") \\\n",
    "        .getOrCreate()\n",
    "\n",
    "def load_data(spark, data_dir):\n",
    "    return spark.read.json(data_dir)\n",
    "\n",
    "def jieba_tokenizer(text):\n",
    "    words = jieba.cut(text)\n",
    "    return [word for word in words]\n",
    "\n",
    "def register_jieba_udf(data):\n",
    "    jieba_udf = udf(jieba_tokenizer, ArrayType(StringType()))\n",
    "    return data.withColumn(\"words\", jieba_udf(data[\"content\"]))\n",
    "\n",
    "\n",
    "\n",
    "# Train and evaluate a logistic regression model\n",
    "def train_and_evaluate_model(train_data, test_data):\n",
    "    # Feature extraction using HashingTF\n",
    "    hashingTF = HashingTF(numFeatures=2**18, inputCol=\"words\", outputCol=\"features\")\n",
    "    train_data = hashingTF.transform(train_data)\n",
    "    test_data = hashingTF.transform(test_data)\n",
    "\n",
    "    # Train a logistic regression model\n",
    "    lr = LogisticRegression(maxIter=10, regParam=0.02, featuresCol=\"features\")\n",
    "    model = lr.fit(train_data)\n",
    "\n",
    "    # Evaluate the model\n",
    "    predictions = model.transform(test_data)\n",
    "    evaluator = BinaryClassificationEvaluator(rawPredictionCol=\"rawPrediction\", labelCol=\"label\")\n",
    "    auc = evaluator.evaluate(predictions)\n",
    "    accuracy = predictions.filter(predictions[\"prediction\"] == predictions[\"label\"]).count() / predictions.count()\n",
    "\n",
    "    return model, auc, accuracy\n",
    "\n",
    "# Save the trained model to disk\n",
    "def save_model(model, model_path):\n",
    "    model.save(model_path)\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    start_time = time.time()\n",
    "    \n",
    "    # Initialize\n",
    "    jieba.initialize() \n",
    "    spark = initialize_spark()\n",
    "\n",
    "    # Load data\n",
    "    data_dir = \"/path/to/data\"\n",
    "    data = load_data(spark, data_dir)\n",
    "\n",
    "    # Split data into train and test sets\n",
    "    train_data, test_data = data.randomSplit([0.8, 0.2], seed=123)\n",
    "\n",
    "    # Register Jieba tokenizer as a UDF\n",
    "    train_data = register_jieba_udf(train_data)\n",
    "    test_data = register_jieba_udf(test_data)\n",
    "\n",
    "    # Train and evaluate the model\n",
    "    model, auc, accuracy = train_and_evaluate_model(train_data, test_data)\n",
    "    print(\"Area Under the ROC Curve (AUC):\", auc)\n",
    "    print(\"Accuracy:\", accuracy)\n",
    "    \n",
    "    # Get size of wights\n",
    "    num_weights = len(model.coefficients)\n",
    "    print(\"Number of Weights (Coefficients) in the Model:\", num_weights)\n",
    "    \n",
    "    # Save the model to disk\n",
    "    model_path = \"/path/to/model\"\n",
    "    save_model(model, model_path)  \n",
    "\n",
    "    end_time = time.time()\n",
    "    execution_time = end_time - start_time\n",
    "    print(\"Code execution time:\", execution_time, \"seconds\")"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
